import numpy as np
import torch
import torch.nn as nn

class DiceLoss(nn.Module):

    def __init__(self):
        super(DiceLoss, self).__init__()

    def forward(self, predicted, label, smooth=1):
        
        # print(predicted.size())
        # print(label.size())

        pred = torch.softmax(predicted, dim=1)

        pred = pred.reshape(-1)
        label = label.reshape(-1)
        
#         class_selectors = tf.cast(K.argmax(label, axis=1), tf.int32)
#         class_selectors = [K.equal(i, class_selectors) for i in range(len(weights_list))]
#         class_selectors = [K.cast(x, K.floatx()) for x in class_selectors]
#         weights = [sel * w for sel, w in zip(class_selectors, weights_list)]
#         weight_multiplier = weights[0]
#         for i in range(1, len(weights)):
#             weight_multiplier = weight_multiplier + weights[i]

        intersection = (pred * label).sum()
        dice = (2.*intersection + smooth)/(pred.sum() + label.sum() + smooth)

#         dice = dice * weight_multiplier
        
        return 1 - dice

class IOULoss(nn.Module):
    def __init___(self):
        super(IOULoss, self).__init__()

    def forward(self, predicted, label, smooth=1):

        pred = torch.softmax(predicted, dim=1)

        pred = pred.reshape(-1)
        label = label.reshape(-1)
        
        intersection = (pred * label).sum()
        dice = (2.*intersection + smooth)/(pred.sum() + label.sum() + smooth)

        iou = 1.0 * dice / (2. - dice)
        
        return 1 - iou